import asyncio
import json
import re
from typing import List, Dict, Optional
from jinja2 import Environment, FileSystemLoader
import json_repair
from research_agent.core.query import Query
from research_agent.core.paths import VERIFY_STATEMENT_CITATION_PROMPT
from research_agent.core.general_llm import LLM
from pyaml_env import parse_config
from research_agent.core.config import Config

prompt_env = Environment()


class StatementCitationVerifier:
    """
    验证语句引用

    该类用于验证输入语句中引用的准确性，通过检索论文相关信息，
    调用大语言模型（LLM）进行验证，并返回验证结果。
    """

    def __init__(self):
        self.query = Query()

        configs = parse_config(Config.YAML_CONFIG)
        self.llm = LLM(config=configs[Config.DEFAULT_MODEL])
        self.prompt_env = Environment(
            loader=FileSystemLoader(VERIFY_STATEMENT_CITATION_PROMPT))
        with open(VERIFY_STATEMENT_CITATION_PROMPT, "r") as f:
            self.verify_statement_citation_prompt_template = prompt_env.from_string(
                f.read())

    def _prepare_messages(self, statement: str, context: str) -> List[Dict]:
        """
        构造发送给 LLM 的消息列表

        根据输入的语句（含引用）和检索到的上下文信息，
        渲染系统和用户角色的提示模板，生成符合 LLM 要求的消息格式。

        Args:
            statement: 包含引用的语句
            context: 检索到的上下文信息

        Returns:
            一个包含系统和用户消息的列表
        """
        system_prompt = self.verify_statement_citation_prompt_template.render(
            role="system")
        user_prompt = self.verify_statement_citation_prompt_template.render(
            role="user", statement_with_citations=statement, retrieved_chunk=context
        )

        return [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]

    async def _process_statement(self, statements_with_citation: dict) -> Dict:
        """
        处理单个语句

        从输入的语句中移除引用标记，提取出引用的语句和引用信息

        Args:
            statements_with_citation: 包含引用的语句

        Returns:
            一个包含引用的语句和引用信息的字典
        """
        statement = re.sub(r"<sup>.*?</sup>", "",
                           statements_with_citation["statement"])
        citation = re.findall(r'<sup>(.*?)</sup>',
                              statements_with_citation["statement"])
        all_citation_list = []

        unique_citations = set()
        for a in citation:
            temp_a = json.loads(json_repair.repair_json(a))
            # 创建唯一标识符，使用论文标题、chunk_id和paper_id的组合
            citation_key = (temp_a["chunk_id"], temp_a["paper_id"])

            # 如果这个引用还没有被添加过
            if citation_key not in unique_citations:
                # 将唯一标识符添加到集合中
                unique_citations.add(citation_key)

                # 创建包含引用信息的字典并添加到列表中
                citation_info = {
                    "chunk_id": temp_a["chunk_id"],
                    "paper_id": temp_a["paper_id"]
                }
                all_citation_list.append(citation_info)
        # 针对每个引用，构造基于 paper_id 的查询请求（top_k 设置为 100）
        paper_queries = [
            self.query.query_by_id(f"{c['paper_id']}", top_k=100) for c in all_citation_list
        ]
        # 并行执行所有论文查询任务
        retrieved_paper_list = await asyncio.gather(*paper_queries)
        retrieved_paper_text = ""

        # 遍历每个引用及其对应的查询结果
        for a_citation, paper_results in zip(all_citation_list, retrieved_paper_list):
            target_chunk_id = a_citation['chunk_id']
            # 遍历该论文的所有 chunk
            for chunk in paper_results:
                # 如果找到匹配的 chunk_id，则格式化并拼接该 chunk 的详细信息
                if str(chunk['chunk_id']) == str(target_chunk_id):
                    retrieved_paper_text += f'''---\n\npaper_title: {chunk['paper_title']}\n\nchunk_id:{chunk['chunk_id']}\n\npaper_id:{chunk['paper_id']}\n\npaper_content:\n\n{chunk['chunk_text']}\n'''

        # 根据原始语句和检索到的论文内容构造 LLM 消息
        verify_statement_citation_prompt = self._prepare_messages(
            statements_with_citation, retrieved_paper_text)

        try:
            # 调用 LLM 完成验证任务，并解析返回的 JSON 数据
            response = await self.llm.completion(verify_statement_citation_prompt)
            response = json_repair.loads(response)
            # 如果返回结果中不包含 "unsupported_papers" 字段，则截取部分上下文后重新调用 LLM
            if "unsupported_papers" not in response:
                verify_statement_citation_prompt = self._prepare_messages(
                    statements_with_citation, retrieved_paper_text[:20000])
                response = await self.llm.completion(verify_statement_citation_prompt)
                response = json_repair.loads(response)
                return response
            return response
        except (json.JSONDecodeError, KeyError) as e:
            # 如果解析 LLM 返回的数据失败，则抛出运行时错误
            raise RuntimeError(f"Failed to parse LLM response: {str(e)}")

    async def verify_statements(self, statements_with_citations: List[Dict]) -> List[Dict]:
        """
        验证一组语句中的引用

        对输入的语句列表逐个进行引用验证，分别返回支持和不支持引用的结果。

        Args:
            statements_with_citations: 包含语句及引用信息的字典列表

        Returns:
            一个元组，包含支持引用的语句列表和不支持引用的语句列表
        """
        if not isinstance(statements_with_citations, list) or not statements_with_citations:
            raise ValueError("Input must be a non-empty list of statements")

        tasks = [self._process_statement(stmt)
                 for stmt in statements_with_citations]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        support_statement_citation_result = []
        not_support_statement_citation_result = []
        for result, stmt in zip(results, statements_with_citations):
            # 如果结果不是字典或者不包含 "unsupported_papers" 字段，则打印并跳过
            if not isinstance(result, dict) or "unsupported_papers" not in result:
                print("*" * 20, result, "*" * 20)
                continue
            unsupport_papers = result["unsupported_papers"]
            # 根据返回的 "unsupported_papers" 字段判断该语句引用是否支持
            if unsupport_papers:
                not_support_statement_citation_result.append(result)
            else:
                support_statement_citation_result.append(stmt)
        # 返回支持和不支持引用的语句列表
        return support_statement_citation_result, not_support_statement_citation_result


async def main():
    verifier = StatementCitationVerifier()
    statements = [
        "The study found that the new drug was effective in treating the disease."
    ]
    results = await verifier.verify_statements(statements)
    print(results)

if __name__ == "__main__":
    asyncio.run(main())
